import time
from typing import Any
import zmq
import threading
import pickle
import os
import numpy as np
from PIL import Image
import cv2
from camera_node import ZMQClientCamera


hostname = "192.168.0.5"
remotename = "192.168.0.4"
# ['317622072614', '317622072616', '317622075898']
# right, front, left
right_wrist_camera_port = 6003
left_wrist_camera_port = 6004
top_camera_port = 6005
arm_recv_port = 6001
base_recv_port = 6002
arm_send_port = 6006
base_send_port = 6007

camera_clients = {
    "left": None,
    "right": None,
    "top": None,
}


class ArxRobot:
    def __init__(self, arm_port = arm_recv_port, base_port = base_recv_port):

        # self.socket_arm_recv, _ = self.create_zmq_connection(socket_type=zmq.SUB,
        #                                                      port=arm_recv_port,
        #                                                      bind=False,
        #                                                      hwm=2,
        #                                                      conflate=True)
        #
        # self.socket_base_recv, _ = self.create_zmq_connection(socket_type=zmq.SUB,
        #                                                       port=base_recv_port,
        #                                                       bind=False,
        #                                                       hwm=2,
        #                                                       conflate=True)


        self.socket_arm_recv = None

        self.socket_base_recv = None

        self.socket_arm_send, _ = self.create_zmq_connection(socket_type=zmq.PUB,
                                                             port=arm_send_port,
                                                             bind=True,
                                                             hwm=2,
                                                             conflate=True)

        self.socket_base_send, _ = self.create_zmq_connection(socket_type=zmq.PUB,
                                                              port=base_send_port,
                                                              bind=True,
                                                              hwm=2,
                                                              conflate=True)

        debug_message = f"Robot Sever Binding to arm {arm_port}, base {base_port} Robot: ARX"
        print(debug_message)
        self._timout_message = f"Timeout in Robot Server, Robot: ARX"
        self._stop_event = threading.Event()

        # self.actions = self.init_data()
        self.home_joints = [0.0 for i in range(14)]

        print("Initializing end")

    def create_zmq_connection(self, socket_type, port, bind=True, context=None, hwm=2, conflate=True):
        # Create new context if none provided
        context = context or zmq.Context()

        socket = context.socket(socket_type)

        if socket_type == zmq.PUB:
            socket.setsockopt(zmq.SNDHWM, hwm)

            # Set socket options based on socket type
            if conflate:
                socket.setsockopt(zmq.CONFLATE, 1)

            address = f"tcp://{hostname}:{port}"
            socket.bind(address)

        elif socket_type == zmq.SUB:
            address = f"tcp://{remotename}:{port}"
            socket.connect(address)

            # Set socket options based on socket type
            if conflate:
                socket.setsockopt(zmq.CONFLATE, 1)

            socket.setsockopt(zmq.RCVTIMEO, 10000) # Set timeout to 1000 ms
            socket.setsockopt(zmq.RCVHWM, hwm)
            socket.setsockopt(zmq.SUBSCRIBE, b'')  # Subscribe to all messages

        return socket, context

    def get_obs(self, request_list):
        self.socket_arm_recv, _ = self.create_zmq_connection(socket_type=zmq.SUB,
                                                             port=arm_recv_port,
                                                             bind=False,
                                                             hwm=2,
                                                             conflate=True)

        self.socket_base_recv, _ = self.create_zmq_connection(socket_type=zmq.SUB,
                                                              port=base_recv_port,
                                                              bind=False,
                                                              hwm=2,
                                                              conflate=True)

        camera_clients = {
            "left": ZMQClientCamera(port=left_wrist_camera_port, host=remotename),
            "right": ZMQClientCamera(port=right_wrist_camera_port, host=remotename),
            "top": ZMQClientCamera(port=top_camera_port, host=remotename),
        }

        """Collects robot observations from specified sensors and cameras."""
        obs = {"images": {}}  # Only initialize images dict by default

        # Configuration for camera parameters and processing
        CAMERA_CONFIG = {
            'cam_left_wrist': {
                'client_key': 'left',
                'resize': None
            },
            'cam_right_wrist': {
                'client_key': 'right',
                'resize': None
            },
            'cam_top': {
                'client_key': 'top',
                'resize': (640, 480)
            }
        }
        
        for name in request_list:
            # Handle camera requests
            if name in CAMERA_CONFIG:
                config = CAMERA_CONFIG[name]

                # Get image from camera client
                image = camera_clients[config['client_key']].read()
                # Convert BGR to RGB format
                image = image[:, :, ::-1]

                # Apply resizing if configured
                if config['resize']:
                    image = Image.fromarray(image).resize(config['resize'])
                    image = np.array(image)

                obs['images'][name] = image
                # print(f"{name.replace('_', ' ')} data received")

            # Handle joint state request
            elif name == 'joints':
                arm_message = self.socket_arm_recv.recv()
                obs['joints'] = pickle.loads(arm_message)

            # Handle base state request
            elif name == 'base':
                base_message = self.socket_base_recv.recv()
                obs['base'] = pickle.loads(base_message)

            # Warn about unrecognized requests
            else:
                print(f"Warning: Unrecognized sensor '{name}' in request list")

        camera_clients['left']._socket.close()
        camera_clients['right']._socket.close()
        camera_clients['top']._socket.close()



        self.socket_arm_recv.close()

        self.socket_base_recv.close()


        return obs

    def step(self, joints, base=None):

        if base is not None:
            if len(base) != 6:
                print(f"base len is {len(base)}! wrong length!")
                return
            if any(np.isnan(x) for x in base):
                print(f"Nan in base!")
                return

            base = [float(x) for x in base]
            self.socket_base_send.send(pickle.dumps(base))
            print(self._timout_message)

        try:
            if len(joints) != 14:
                print(f"joints len is {len(joints)}! wrong length!")
                return
            if any(np.isnan(x) for x in joints):
                print(f"Nan in joints!")
                return

            joints = [float(x) for x in joints]
            print(f"get the joint command {joints}")
            self.socket_arm_send.send(pickle.dumps(joints))
        except zmq.Again:
            print(self._timout_message)

    def stop(self) -> None:
        """Signal the server to stop serving."""
        self._stop_event.set()

    def reset(self, request2reset=['joints']):
        '''reset to init pos'''
        time_step = 0.1
        step = np.deg2rad(90) * 0.02
        base_reset = True
        joints_reset = True
        base_home = [0.0] * 6
        joints_home = [0.0] * 14
        def ramp(goal, current, ramp_k):
            delta = goal - current
            # Element-wise check: set to goal if within ramp_k, otherwise step towards goal
            return np.where(
                np.abs(delta) <= ramp_k,  # Condition per joint
                goal,  # Set to goal if condition met
                current + np.sign(delta) * ramp_k  # Step towards goal otherwise
            )

        if "joints" in request2reset:
            joints_reset = False

        if "base" in request2reset:
            base_reset = False

        while not base_reset or not joints_reset:
            obs = self.get_obs(request2reset)
            if not base_reset:
                now_base = obs['base']
                if np.allclose(now_base, base_home, atol=0.01):
                    base_reset = True
                self.socket_base_send.send(pickle.dumps(base_home))

            if not joints_reset:
                now_joints = obs['joints']
                now_joints[6] = (now_joints[6]-0)/5
                now_joints[13] = (now_joints[13]-0)/5
                print(now_joints)
                if np.allclose(now_joints, joints_home, atol=step * 2):
                    joints_reset = True
                cmds = ramp(joints_home, now_joints, step)
                self.socket_arm_send.send(pickle.dumps(cmds))

            time.sleep(0.02)
            pass

def main():
    robot = ArxRobot()
    time.sleep(2)
    robot.reset()

    # obs = robot.get_obs(['cam_left_wrist', 'cam_right_wrist', 'cam_top', 'joints'])
    # print("start")
    # print(obs['joints'])
    # start = time.time()
    # # while time.time() - start < 5:
    # #     pass
    # obs = robot.get_obs(['cam_left_wrist', 'cam_right_wrist', 'cam_top', 'joints'])
    # print(obs['joints'])
    # # print("AAAA")
    # #
    # while time.time() - start < 1:
    #     pass
    # obs = robot.get_obs(['cam_left_wrist', 'cam_right_wrist', 'cam_top', 'joints'])
    # print(obs['joints'])
    # print("BBBB")

    while True:
    #
        #obs = robot.get_obs(['cam_right_wrist','joints'])
        obs = robot.get_obs(['cam_left_wrist', 'cam_right_wrist', 'cam_top', 'joints'])
        #obs = robot.get_obs(['cam_left_wrist', 'cam_right_wrist', 'cam_top'])
        #obs = robot.get_obs(['joints'])

        #print(obs['joints'])
        robot.step([0.0,0.0,0.0,0.0,0.0,0.0,5.0,0.0,0.0,0.0,0.0,0.0,0.0,5.0])
        array_top = obs['images']['cam_top']
        array_right = obs['images']['cam_right_wrist']
        array_left = obs['images']['cam_left_wrist']
        # print("!")
        cv2.imshow('camera', array_top)
        cv2.imshow('camera2', array_right)
        cv2.imshow('camera3', array_left)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

        time.sleep(0.02)

if __name__ == '__main__':
    main()
